import functools
import math
import multiprocessing
import random
import signal
import sys

import nfa
import cfg
import utils

def sample_fraction(l, fraction):
  return random.sample(l, int(len(l) / fraction))

def subprocess_init(machine_json):
  signal.signal(signal.SIGINT, signal.SIG_IGN)
  global machine
  machine = nfa.from_json(machine_json)

def subprocess_task(word):
  return machine.parse(word)

class EnergyComputer:
  def __init__(self, terminals):
    self.trivial_dl_per_symbol = utils.bits_required(len(terminals))
  def compute(self, machine, words):
    with multiprocessing.Pool(initializer=subprocess_init,
                              initargs=[machine.to_json()]) as pool:
      dls = pool.map(subprocess_task, words)
      for i in range(len(dls)):
        if not dls[i]:
          # dls[i] = len(self.words[i]) * self.trivial_dl_per_symbol
          return False
      return sum(dls) + machine.description_length()

class Experiment:
  def __init__(self, filename=None):
    self.logger = utils.Logger(filename)
  def decide(self, old_energy, new_energy, temperature, greedy=True):
    self.logger.log('Old energy is {:.2f} and new is {:.2f}'.format(old_energy, new_energy))
    if greedy and old_energy > new_energy:
      return True
    delta = new_energy - old_energy
    probability = 1 / (1 + math.exp(- delta / temperature))
    self.logger.log('Probability to move to neighbor: {:.2%}'.format(1 - probability))
    return random.random() > probability
  def start(self, words, terminals, temperature=500, iterations=5000,
            cfg_or_nfa=False, machine=None, fraction=1):
    delta_t = temperature / iterations
    if cfg_or_nfa:
      if not machine:
        machine = cfg.create_initial_cfg(terminals=terminals)
    else:
      if not machine:
        machine = nfa.create_initial_nfa(terminals=terminals)
    machine_stack = [] # contains all machines reached in the process
    energy_computer = EnergyComputer(terminals)
    trivial_energy = sum(len(word) for word in words) * utils.bits_required(len(terminals))
    old_energy = energy_computer.compute(machine, words)
    self.logger.log('Starting energy is {:.2f}.'.format(old_energy))
    elapsed = 0
    while iterations > 0:
      self.logger.log()
      with utils.Timer():
        self.logger.log('Iteration {} with temperature {:.2f}'.format(iterations, temperature))
        while True:
          neighbor, message = machine.get_neighbor()
          self.logger.log(message)
          if fraction > 1:  
            sample = sample_fraction(words, fraction)
            while True:
              old_energy = energy_computer.compute(machine, sample)
              if old_energy:
                machine_stack.append(machine)
                break
              else:
                self.logger.log('Current machine can\'t parse all words, rolling back once.')
                machine = machine_stack.pop()
            self.logger.log('Current energy recomputed per new sampling: {:.2f} ({:.2%})'.format(old_energy, old_energy / trivial_energy), utils.Logger.OKBLUE)
          else:
            sample = words
          new_energy = energy_computer.compute(neighbor, sample)
          if new_energy:
            break
          else:
            self.logger.log('Neighbor cannot parse all words, getting new neighbor.')
        if self.decide(old_energy, new_energy, temperature):
          self.logger.log('Decision: moved')
          machine = neighbor
          old_energy = new_energy
          self.logger.log('New energy is {:.2f} ({:.2%}) of which {:.2f} for the machine description)'.format(new_energy, new_energy / trivial_energy, machine.description_length()))
          self.logger.log(machine.to_json())
        else:
          self.logger.log('Decision: stayed')
        iterations -= 1
        elapsed += 1
        temperature -= delta_t
  def start_with_data(self, machine_dict=None, iterations=10000,
                      temperature=1000, fraction=1, path='fasta/domains.fasta'):
    self.logger.log('Path for data: {}\n'.format(path))
    machine = None
    if machine_dict:
      import json
      machine = nfa.from_json(json.dumps(machine_dict))
    if path == 'fasta/SCOP_MLA18SS.fasta':
      amino_acids, sequences = read_MLA18SS(path)
    else:
      amino_acids, sequences = read_fasta(path)
    self.start(words=list(sequences.values()), terminals=amino_acids,
               machine=machine, iterations=iterations, temperature=temperature,
               fraction=fraction)

def read_fasta(path):
  def trunc_at(s, b, start ,end):
    return b.join(s.split(b)[start:end])
  protein_sequences = {}
  terminals = set()
  f = open(path, 'r')
  sequences = f.read().split('\n>')
  c = 0
  bad_sequences = 0
  for entry in sequences[1:]:
    c += 1
    rows = entry.splitlines()
    name = rows[0]
    sequence = ''
    for item in rows[1:]:
      sequence += item
    sequence = sequence.lower()
    if (len(sequence) > 0):
      if '-' not in sequence and 'x' not in sequence:
        terminals |= set(list(sequence))
        protein_sequences[name] = sequence
    else:
      bad_sequences += 1
  print('number of protein sequences in fasta file: {}'.format(c))
  print('{} were ignored due to bad sequences.'.format(bad_sequences))
  print('Set of terminals: {}'.format(terminals))
  return list(terminals), protein_sequences

def read_MLA18SS(path='fasta/SCOP_MLA18SS.fasta', use_allowed_only=True):
  _, sequences = read_fasta(path)
  protein_names = list(sequences.keys())
  if use_allowed_only:
    f = open('fasta/SCOP_good_domains_list.txt', 'r')
    allowed = f.read().split('\n')
    protein_names = [x for x in protein_names if x in allowed]
  print('Reading and processing MLA18SS data.')
  terminals = set()
  result = {}
  for k in protein_names:
    # turn 'other' to 'polar' s->i t->q w->l
    result[k] = sequences[k].replace('s', 'i').replace('t', 'q').replace('w', 'l')
    # collect terminals
    terminals |= set(list(result[k]))
  print('Set of terminals after processing: {}'.format(terminals))
  print('Number of protein sequences: {}'.format(len(result)))
  return list(terminals), result

def combine_and_evaluate(a, b, words):
  import json
  c = nfa.from_json(json.dumps(a)) | nfa.from_json(json.dumps(b))
  evaluate_nfa(machine=c, words=words, terminals=c.terminals)

def evaluate_nfa(machine, words, fraction=1):
  sampling = sample_fraction(words, fraction)
  energy_computer = EnergyComputer(machine.terminals)
  print(energy_computer.compute(machine, sampling))

if __name__ == '__main__':
  amino_acids, sequences = read_fasta('fasta/domains.fasta')
  exp = Experiment()
  exp.start(words=list(sequences.values()), terminals=amino_acids)
